import r2pipe
import sys
import re
import subprocess as psutil
import shutil
from itertools import takewhile, dropwhile
addr = 18
dll_file = './ida/cpettemp.dll'
dump_bin = 'E:/code/msvc/VC/Tools/MSVC/14.41.34120/bin/Hostx64/x64/dumpbin.exe'
if len(sys.argv)>1:
    dll_file = sys.argv[1]
if len(sys.argv)>2:
    addr = int(sys.argv[2])
print(dll_file)
shutil.copy2(dll_file, 'shrink.dll')
ps = psutil.Popen([dump_bin, '/DISASM', dll_file], stdout=psutil.PIPE)
with open('file_addr.txt', 'wt') as fp:
    for line in psutil.check_output(['grep', '\[[0]\+18[0-9A-Fa-f]\+h\]'], stdin=ps.stdout).decode().split('\n'):
        try:    
            line = line.strip()
            rip, _, op_code = line.partition(':')
            rip = int(rip, base=16)
            op_code = re.split(r'[ ]{1,}', op_code.strip())
            byte_code = ''.join(takewhile(lambda x:re.match(r'^[0-9a-fA-F]{2}$', x), op_code))
            asm_code = ' '.join(dropwhile(lambda x:re.match(r'^[0-9a-fA-F]{2}$', x), op_code))
            addr =   ''.join(takewhile(lambda x:x!=']', (dropwhile(lambda x:x!='[', asm_code))))[1:-1]
            addr = int(addr, base=16)
            n_rip = rip + len(byte_code)//2
            if (addr-n_rip) <= 0:
                continue
            offset = hex(addr-n_rip)[-4:-2].replace('x', '0').upper()
            
            changed_code =  \
             ''.join([hex(int(offset,base=16)-0x10)[2:].zfill(2) if offset[0]==v1 and offset[1]==v2 else f'{v1}{v2}' for v1, v2 in zip(byte_code[::2], byte_code[1::2])])

            file_offset = rip-0x180001000+0x400
            print(hex(rip), hex(file_offset), byte_code, changed_code, hex(addr-n_rip), offset, hex(addr), asm_code)
            print(hex(file_offset), changed_code, file=fp)
            
        except:
            pass

with open('shrink.dll', 'r+b') as wfp,  open('file_addr.txt', 'rt') as fp:
    for line in fp:
        addr,_, op_code = line.partition(' ')
        op_code = op_code.strip('\n')
        wfp.seek(int(addr, base=16))
        for i in range(0, len(op_code), 2):
            try:
                wfp.write(int(op_code[i:i+2], base=16).to_bytes())
            except Exception as e:
                print(addr, op_code)
            #wfp.seek(1, 1)

"""
r2 = r2pipe.open(dll_file)
for line in r2.cmd(f's 0x180001120;pd 500000~&[0x{addr}').split('\n'):
    if line:
        items = re.split(r'[ ]{2,}', line.strip().partition(';')[0].strip())
        rip = int(items[0], base=16)+len(items[1])//2
        addr =   ''.join(takewhile(lambda x:x!=']', (dropwhile(lambda x:x!='[', items[2]))))[1:]
        op_code = items[1]
        if addr.startswith('0x'):
            try:
                addr = int(addr, base=16)
                print(hex(rip), op_code, hex(int(op_code, base=16)-0x100000), hex(addr), hex(addr-rip), items[2])
            except:
                pass
"""
